//
//  CTRSAUtil.m
//  CTOpenSSL
//
//  Created by Andney Jack on 2021/2/22.
//

#import "CTRSAUtil.h"
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/md5.h>

@interface CTRSAUtil ()
{
    RSA *_public_rsa;
    RSA *_private_rsa;
    BOOL _isRSA2;
}
@end

@implementation CTRSAUtil

- (instancetype)init {
    self = [super init];
    if (self) {
        _public_rsa = NULL;
        _private_rsa = NULL;
        _isRSA2 = NO;
    }
    return self;
}

- (instancetype)initWithRSA2:(BOOL)rsa2 {
    self = [super init];
    if (self) {
        _public_rsa = NULL;
        _private_rsa = NULL;
        _isRSA2 = rsa2;
    }
    return self;
}

#pragma mark - —————— Public Method ——————

- (BOOL)importKey:(NSString *)key keyType:(RSAKeyType)keyType
{
    if (!key || key.length < 5) {
        return NO;
    }

    BOOL status = NO;
    RSA *rsa = NULL;
    BIO *bio = BIO_new(BIO_s_file());

    NSString *rsaFilePath = [NSTemporaryDirectory() stringByAppendingPathComponent:@"tmp_rsa_key"];
    NSString *formattedKey = [self.class formatRSAKey:key keyType:keyType];
    BOOL writeSuccess = [formattedKey writeToFile:rsaFilePath atomically:YES encoding:NSUTF8StringEncoding error:nil];
    if (!writeSuccess) {
        return NO;
    }

    const char *cPath = [rsaFilePath cStringUsingEncoding:NSUTF8StringEncoding];
    BIO_read_filename(bio, cPath);
    if (keyType == RSAKeyTypePrivate) {
        rsa = PEM_read_bio_RSAPrivateKey(bio, NULL, NULL, "");
        _private_rsa = rsa;
        if (rsa != NULL && 1 == RSA_check_key(rsa)) {
            status = YES;
        } else {
            status = NO;
        }
    } else {
        rsa = PEM_read_bio_RSA_PUBKEY(bio, NULL, NULL, NULL);
        _public_rsa = rsa;
        if (rsa != NULL) {
            status = YES;
        } else {
            status = NO;
        }
    }

    BIO_free_all(bio);
    [[NSFileManager defaultManager] removeItemAtPath:rsaFilePath error:nil];
    return status;
}

- (BOOL)importKeyWithPath:(NSString *)path keyType:(RSAKeyType)type
{
    BOOL status = NO;
    const char *cPath = [path cStringUsingEncoding:NSUTF8StringEncoding];
    FILE *file = fopen(cPath, "rb");
    if (!file) {
        return status;
    }

    if (type == RSAKeyTypePublic) {
        _public_rsa = NULL;
        if ((_public_rsa = PEM_read_RSA_PUBKEY(file, NULL, NULL, NULL))) {
            status = YES;
        }
    } else if (type == RSAKeyTypePrivate) {
        _private_rsa = NULL;
        if ((_private_rsa = PEM_read_RSAPrivateKey(file, NULL, NULL, NULL))) {
            status = YES;
        }
    }
    fclose(file);
    return status;
}

+ (int)getBlockSize:(RSA *)rsa paddingType:(int)type
{
    int len = RSA_size(rsa);
    if (type == RSA_PKCS1_PADDING || type == RSA_SSLV23_PADDING) {
        len -= 11;
    } else if (type == RSA_X931_PADDING) {
        len -= 2;
    }
    return len;
}

+ (NSString *)formatRSAKey:(NSString *)key keyType:(RSAKeyType)type
{
    NSInteger lineNum = -1;
    NSMutableString *result = [NSMutableString string];

    if (type == RSAKeyTypePrivate) {
        [result appendString:@"-----BEGIN PRIVATE KEY-----\n"];
        lineNum = 79;
    } else if (type == RSAKeyTypePublic) {
        [result appendString:@"-----BEGIN PUBLIC KEY-----\n"];
        lineNum = 76;
    }

    int count = 0;
    for (int i = 0; i < [key length]; ++i) {
        unichar c = [key characterAtIndex:i];
        if (c == '\n' || c == '\r') {
            continue;
        }
        [result appendFormat:@"%c", c];
        if (++count == lineNum) {
            [result appendString:@"\n"];
            count = 0;
        }
    }

    if (type == RSAKeyTypePrivate) {
        [result appendString:@"\n-----END PRIVATE KEY-----"];
    } else if (type == RSAKeyTypePublic) {
        [result appendString:@"\n-----END PUBLIC KEY-----"];
    }
    return result;
}

/// RSA sha1验证签名
// sign为Base64字符串
- (BOOL)verifyString:(NSString *)string withSign:(NSString *)sign keyType:(RSAKeyType)keyType
{
    if (![self m_checkRSAExist:keyType]) {
        return NO;
    }

    const char *message = [string cStringUsingEncoding:NSUTF8StringEncoding];
    NSUInteger messageLength = [string lengthOfBytesUsingEncoding:NSUTF8StringEncoding];
    NSData *signatureData = [[NSData alloc] initWithBase64EncodedString:sign options:NSDataBase64DecodingIgnoreUnknownCharacters];
    unsigned char *sig = (unsigned char *)[signatureData bytes];
    unsigned int sig_len = (int)[signatureData length];

    int bufferLength = _isRSA2 ? SHA256_DIGEST_LENGTH : SHA_DIGEST_LENGTH;
    unsigned char buffer[bufferLength];

    if (_isRSA2) {
        SHA256((unsigned char *)message, messageLength, buffer);
    } else {
        SHA1((unsigned char *)message, messageLength, buffer);
    }

    int type = _isRSA2 ? NID_sha256 : NID_sha1;
    RSA *rsa = (keyType == RSAKeyTypePublic) ? _public_rsa : _private_rsa;
    int verify_ok = RSA_verify(type, buffer, bufferLength
                               , sig, sig_len
                               , rsa);

    if (1 == verify_ok) {
        return YES;
    }
    return NO;
}

#pragma mark RSA MD5 验证签名

- (BOOL)verifyMD5String:(NSString *)string withSign:(NSString *)sign keyType:(RSAKeyType)keyType
{
    if (![self m_checkRSAExist:keyType]) {
        return NO;
    }

    const char *message = [string cStringUsingEncoding:NSUTF8StringEncoding];
    NSData *signatureData = [[NSData alloc] initWithBase64EncodedString:sign options:NSDataBase64DecodingIgnoreUnknownCharacters];
    unsigned char *sig = (unsigned char *)[signatureData bytes];
    unsigned int sig_len = (int)[signatureData length];

    unsigned char digest[MD5_DIGEST_LENGTH];
    MD5_CTX ctx;
    MD5_Init(&ctx);
    MD5_Update(&ctx, message, strlen(message));
    MD5_Final(digest, &ctx);

    RSA *rsa = (keyType == RSAKeyTypePublic) ? _public_rsa : _private_rsa;
    int verify_ok = RSA_verify(NID_md5
                               , digest, MD5_DIGEST_LENGTH
                               , sig, sig_len
                               , rsa);
    if (1 == verify_ok) {
        return YES;
    }
    return NO;
}

#pragma mark - ————————————————

- (NSString *)signString:(NSString *)string keyType:(RSAKeyType)keyType
{
    return [self signString:string keyType:keyType urlEncoding:NO];
}

- (NSString *)signString:(NSString *)string keyType:(RSAKeyType)keyType urlEncoding:(BOOL)urlEncoding
{
    if (![self m_checkRSAExist:keyType]) {
        return nil;
    }

    const char *message = [string cStringUsingEncoding:NSUTF8StringEncoding];
    int messageLength = (int)strlen(message);
    unsigned char *sig = (unsigned char *)malloc(256);
    unsigned int sig_len;

    int bufferLength = _isRSA2 ? SHA256_DIGEST_LENGTH : SHA_DIGEST_LENGTH;
    unsigned char sha_buffer[bufferLength];
    if (_isRSA2) {
        SHA256((unsigned char *)message, messageLength, sha_buffer);
    } else {
        SHA1((unsigned char *)message, messageLength, sha_buffer);
    }

    int type = _isRSA2 ? NID_sha256 : NID_sha1;
    RSA *rsa = (keyType == RSAKeyTypePublic) ? _public_rsa : _private_rsa;
    int rsa_sign_valid = RSA_sign(type
                                  , sha_buffer, bufferLength
                                  , sig, &sig_len
                                  , rsa);
    if (rsa_sign_valid == 1) {
        NSData *data = [NSData dataWithBytes:sig length:sig_len];

        NSString *base64String = [data base64EncodedStringWithOptions:0];
        free(sig);
        if (urlEncoding) {
            return [self.class urlEncode:base64String];
        } else {
            return base64String;
        }
    }

    free(sig);
    return nil;
}

- (NSString *)signMD5String:(NSString *)string keyType:(RSAKeyType)keyType
{
    if (![self m_checkRSAExist:keyType]) {
        return nil;
    }

    const char *message = [string cStringUsingEncoding:NSUTF8StringEncoding];
    unsigned char *sig = (unsigned char *)malloc(256);
    unsigned int sig_len;

    unsigned char digest[MD5_DIGEST_LENGTH];
    MD5_CTX ctx;
    MD5_Init(&ctx);
    MD5_Update(&ctx, message, strlen(message));
    MD5_Final(digest, &ctx);

    RSA *rsa = (keyType == RSAKeyTypePublic) ? _public_rsa : _private_rsa;
    int rsa_sign_valid = RSA_sign(NID_md5
                                  , digest, MD5_DIGEST_LENGTH
                                  , sig, &sig_len
                                  , rsa);

    if (rsa_sign_valid == 1) {
        NSData *data = [NSData dataWithBytes:sig length:sig_len];

        NSString *base64String = [data base64EncodedStringWithOptions:0];
        free(sig);
        return base64String;
    }

    free(sig);
    return nil;
}

- (NSString *)encode:(NSString *)content keyType:(RSAKeyType)keyType urlEncoding:(BOOL)urlEncoding
{
    if (![self m_checkRSAExist:keyType]) {
        return nil;
    }

    NSString *srcText = urlEncoding ? [self.class urlEncode:content] : [content copy];

    int status;
    int length = (int)[srcText length];
    unsigned char input[length + 1];
    bzero(input, length + 1);
    for (int i = 0; i < length; i++) {
        input[i] = [srcText characterAtIndex:i];
    }

    RSA *rsa = (keyType == RSAKeyTypePublic) ? _public_rsa : _private_rsa;
    int paddingType = RSA_PKCS1_PADDING;

    NSInteger flen = [self.class getBlockSize:rsa paddingType:paddingType];

    char *encData = (char *)malloc(flen);
    bzero(encData, flen);
    if (keyType == RSAKeyTypePublic) {
        status = RSA_public_encrypt(length, (unsigned char *)input, (unsigned char *)encData, rsa, paddingType);
    } else {
        status = RSA_private_encrypt(length, (unsigned char *)input, (unsigned char *)encData, rsa, paddingType);
    }

    if (status) {
        NSData *returnData = [NSData dataWithBytes:encData length:status];
        free(encData);
        encData = NULL;

        NSString *ret = [returnData base64EncodedStringWithOptions:NSDataBase64Encoding64CharacterLineLength];
        return ret;
    }

    free(encData);
    encData = NULL;

    return nil;
}

- (NSString *)decode:(NSString *)content keyType:(RSAKeyType)keyType urlDecoding:(BOOL)urlDecoding
{
    NSData *data = [[NSData alloc] initWithBase64EncodedString:content options:NSDataBase64DecodingIgnoreUnknownCharacters];
    NSData *resultData = [self m_decodeData:data keyType:keyType];
    if (resultData) {
        NSString *decryptString = [[NSString alloc] initWithData:resultData encoding:NSASCIIStringEncoding];
        if (urlDecoding) {
            return [self.class urlDecode:decryptString];
        } else {
            return decryptString;
        }
    }
    return nil;
}

- (NSData *)m_decodeData:(NSData *)data keyType:(RSAKeyType)keyType
{
    if (![self m_checkRSAExist:keyType]) {
        return nil;
    }

    int status;

    int length = (int)[data length];

    RSA *rsa = (keyType == RSAKeyTypePublic) ? _public_rsa : _private_rsa;
    int paddingType = RSA_PKCS1_PADDING;

    NSInteger flen = [self.class getBlockSize:rsa paddingType:paddingType];
    char *decData = (char *)malloc(flen);
    bzero(decData, flen);

    if (keyType == RSAKeyTypePublic) {
        status = RSA_public_decrypt(length, (unsigned char *)[data bytes], (unsigned char *)decData, _public_rsa, paddingType);
    } else {
        status = RSA_private_decrypt(length, (unsigned char *)[data bytes], (unsigned char *)decData, _private_rsa, paddingType);
    }

    if (status) {
        NSData *resultData = [[NSData alloc] initWithBytes:decData length:strlen(decData)];

        free(decData);
        decData = NULL;

        return resultData;
    }

    free(decData);
    decData = NULL;

    return nil;
}

- (NSString *)encodeLongText:(NSString *)content keyType:(RSAKeyType)keyType urlEncoding:(BOOL)urlEncoding
{
    if (![self m_checkRSAExist:keyType]) {
        return nil;
    }

    NSString *srcText = urlEncoding ? [self.class urlEncode:content] : [content copy];
    int maxLength = _isRSA2 ? 245 : 117;
    maxLength = 117;

    NSMutableString *result = [[NSMutableString alloc] init];
    for (NSInteger i = 0; i < ceilf(srcText.length * 1.0 / maxLength); i++) {
        NSString *subStr = [srcText substringWithRange:NSMakeRange(i * maxLength, MIN(maxLength, srcText.length - i * maxLength))];
        // 不能再次进行URL编码
        NSString *subEncodedText = [self encode:subStr keyType:keyType urlEncoding:NO];
        if (subEncodedText) {
            [result appendString:subEncodedText];
        }
    }

    return result.length > 0 ? result : nil;
}

- (NSString *)decodeLongText:(NSString *)content keyType:(RSAKeyType)keyType urlDecoding:(BOOL)urlDecoding
{
    if (![self m_checkRSAExist:keyType]) {
        return nil;
    }

    NSData *data = [[NSData alloc] initWithBase64EncodedString:content options:NSDataBase64DecodingIgnoreUnknownCharacters];

//    int maxLength = _isRSA2 ? 245 : 117;
    int maxLength = _isRSA2 ? 256 : 128;
    maxLength = 128;

    NSMutableData *decryptData = [[NSMutableData alloc] init];

    NSInteger maxTime = ceilf(data.length * 1.0 / maxLength);
    for (NSInteger i = 0; i < maxTime; i++) {
        NSData *subData = [data subdataWithRange:NSMakeRange(i * maxLength, MIN(maxLength, data.length - i * maxLength))];
        NSData *resultData = [self m_decodeData:subData keyType:keyType];
        if (resultData) {
            [decryptData appendData:resultData];
        }
    }

    if (decryptData.length > 0) {
        NSString *decryptString = [[NSString alloc] initWithData:decryptData encoding:NSASCIIStringEncoding];
        if (urlDecoding) {
            return [self.class urlDecode:decryptString];
        } else {
            return decryptString;
        }
    }

    return nil;
}

- (BOOL)m_checkRSAExist:(RSAKeyType)keyType
{
    if (keyType == RSAKeyTypePublic) {
        if (!_public_rsa) {
#ifdef DEBUG
            NSLog(@"please import public key first");
#endif
            return NO;
        }
    } else if (keyType == RSAKeyTypePrivate) {
        if (!_private_rsa) {
#ifdef DEBUG
            NSLog(@"please import private key first");
#endif
            return NO;
        }
    } else {
        return NO;
    }

    return YES;
}

+ (NSString *)urlEncode:(NSString *)text
{
    NSCharacterSet *charset = [[NSCharacterSet characterSetWithCharactersInString:@"!*'();:@&=+$,/?%#[]"] invertedSet];
    return [text stringByAddingPercentEncodingWithAllowedCharacters:charset];
}

+ (NSString *)urlDecode:(NSString *)text
{
    return [text stringByRemovingPercentEncoding];
}

@end
